from logging import PlaceHolder
import streamlit as st
import requests
import json
import uuid

# 确保streamlit_option_menu已安装
# 如果没有安装，可以通过pip install streamlit-option-menu安装

def init_ui():
    # 页面配置
    # read assets/icon.ico as the page_icon
    with open("assets/icon.ico", "rb") as f:
        icon = f.read()
    st.set_page_config(
        page_title="SigOne Web Chat",
        page_icon=icon,
        layout="centered",
    )
    
    # 初始化会话状态
    if "server_url" not in st.session_state:
        st.session_state.server_url = "http://localhost:7891/api/chat/stream"
    if "model_name" not in st.session_state:
        st.session_state.model_name = "qwen2.5-vl-32b"
    if "max_plan_iterations" not in st.session_state:
        st.session_state.max_plan_iterations = "1"
    if "max_step_num" not in st.session_state:
        st.session_state.max_step_num = "5"
    if "max_search_results" not in st.session_state:
        st.session_state.max_search_results = "3"
    if "show_intermediate_results" not in st.session_state:
        st.session_state.show_intermediate_results = False

    # set title and align it to the center
    st.markdown("""
    <style>
        .title {
            text-align: center;
            font-size: 3em;
            font-weight: bold;
        }
    </style>
    <div class="title">SigOne Web Chat</div>
    """, unsafe_allow_html=True)
    chat, settings = st.tabs(["Chat", "Settings"])
    with settings:
        st.text_input("Model", key="model_name", value=st.session_state.model_name,
                      on_change=lambda: setattr(st.session_state, 
                                                "model_name", 
                                                st.session_state.model_name))
        st.text_input("Agent Server URL", key="server_url", value=st.session_state.server_url,
                      on_change=lambda: setattr(st.session_state, 
                                                "server_url",
                                                st.session_state.server_url))
        st.text_input("Max Plan Iterations", key="max_plan_iterations", value=st.session_state.max_plan_iterations,
                      on_change=lambda: setattr(st.session_state, 
                                                "max_plan_iterations",
                                                st.session_state.max_plan_iterations))
        st.text_input("Max Step Num", key="max_step_num", value=st.session_state.max_step_num,
                      on_change=lambda: setattr(st.session_state, 
                                                "max_step_num",
                                                st.session_state.max_step_num))
        st.text_input("Max Search Results", key="max_search_results", value=st.session_state.max_search_results,
                      on_change=lambda: setattr(st.session_state, 
                                                "max_search_results",
                                                st.session_state.max_search_results))
        st.checkbox("Show Intermediate Results", key="show_intermediate_results",
                    value=st.session_state.show_intermediate_results,
                    on_change=lambda: setattr(st.session_state, 
                                              "show_intermediate_results",
                                              st.session_state.show_intermediate_results))
        if st.button("🗑️ Clear History"):
            reset_history()
            st.rerun()
    
    # 创建聊天历史容器
    st.session_state.chat_container = chat
    
    return False, st.session_state

# 会话管理模块

def init_session():
    if "messages" not in st.session_state:
        st.session_state.messages = [
            {"role": "assistant", "content": "How can I help you today?"}
        ]

def add_user_message(content):
    st.session_state.messages.append({"role": "user", "content": content})

def add_assistant_message(content):
    st.session_state.messages.append({"role": "assistant", "content": content})

def display_chat_history():
    with st.session_state.chat_container:
        for msg in st.session_state.messages:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])

def reset_history():
    st.session_state.messages = [
        {"role": "assistant", "content": "How can I help you today?"}
    ]

# API通信模块
def call_agent_server(server_url, model, max_plan_iterations, max_step_num, max_search_results, messages):
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json"
    }
    
    thread_id = str(uuid.uuid4())
    payload = {
        "thread_id": thread_id,
        "model": model,
        "max_plan_iterations": max_plan_iterations,
        "max_step_num": max_step_num,
        "max_search_results": max_search_results,
        "messages": messages
    }
    
    # 创建中间结果展示框
    if st.session_state.show_intermediate_results:
        intermediate_container = st.container(border=True)
    
    try:
        response = requests.post(
            server_url,
            headers=headers,
            data=json.dumps(payload),
            stream=True,
            timeout=500  # 500秒超时
        )
        
        if response.status_code != 200:
            return {
                "error": f"API Error {response.status_code}",
                "details": response.text
            }
        
        full_response = ""
        
        intermediate_results = ""
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')
                
                if decoded_line.startswith("event:"):
                    event_type = decoded_line[6:].strip()
                    intermediate_results = f"event_type: {event_type}"
                    continue
                elif decoded_line.startswith("data:"):
                    decoded_line = decoded_line[5:].strip()
                    intermediate_results = f"data: {decoded_line}"

                try:
                    response_text = json.loads(decoded_line)
                    
                    if isinstance(response_text, dict):
                        if 'langgraph_node' in response_text and response_text['langgraph_node'] in ['planner', 'background_investigator']:
                            continue

                        if 'content' in response_text:
                            full_response += response_text["content"]

                except json.JSONDecodeError as e:
                    error_msg = f"JSON解析错误: {str(e)}"
                    full_response += error_msg
                
        return full_response
            
    except Exception as e:
        error_msg = f"连接失败: {str(e)}"
        return error_msg

# 响应处理模块

def handle_response(api_response):
    if "error" in api_response:
        return f"Error: {api_response.get('error', 'Unknown error')}\n" \
               f"Details: {api_response.get('details', 'No details')}"
    
    
    
    return api_response

# 主应用逻辑

def main():
    # 初始化UI和会话
    reset, session_state = init_ui()
    init_session()
    
    # 显示历史消息
    display_chat_history()
    
    # 处理重置请求
    if reset:
        reset_history()
        st.rerun()
    
    # 处理用户输入
    if user_input := st.chat_input("Input your question?"):
        add_user_message(user_input)

        # 显示用户消息
        with st.session_state.chat_container:
            with st.chat_message("user"):
                st.write(user_input)

            # 显示加载状态
            with st.spinner("Thinking...", show_time=True):
                with st.chat_message("assistant"):
                    api_response = call_agent_server(
                        st.session_state.server_url,
                        st.session_state.model_name,
                        int(st.session_state.max_plan_iterations),
                        int(st.session_state.max_step_num),
                        int(st.session_state.max_search_results),
                        st.session_state.messages
                    )

                    st.markdown(api_response)

                # 处理响应
                # assistant_response = handle_response(api_response)

            # with st.chat_message("assistant"):
            #     st.write(api_response)

        # 添加助手消息到会话历史
        add_assistant_message(api_response)

        # 清空输入框
        st.session_state.user_input = ""
        # st.rerun()

if __name__ == "__main__":
    main()
